#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import random
import argparse
import os
import numpy as np

import cv2

# SLEEP_TIME = 40
SLEEP_TIME = 1

# base_path = r'/rootfs/media/kasim/DataSet/Record'
# out_path = r'/rootfs/media/kasim/DataSet/Record'
# video_file_list = r'/rootfs/media/kasim/DataSet/Record/video_attr_filter_file_list.txt'

base_path = r'/rootfs/media/kasim/DataSet/RecordEx'
out_path = r'/rootfs/media/kasim/DataSet/RecordEx'
video_file_list = r'/rootfs/media/kasim/DataSet/RecordEx/video_attr_filter_file_list.txt'

DISPLAY_WIDTH = 1280
DISPLAY_HEIGHT = 720


def draw_frame(image, frame_num, total_frames_num):

    text_scale = max(1.0, image.shape[1] / 800.)

    cv2.putText(image, 'frame: %d/%d' % (frame_num, total_frames_num),
                (0, int(15 * text_scale)), cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 255), thickness=2)


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    filter_select_video_paths = set()
    select_file_path = os.path.join(out_path, 'video.select.txt')
    if os.path.exists(select_file_path):
        with open(select_file_path, 'r') as file:
            for video_path in file.readlines():
                filter_select_video_paths.add(video_path.strip())

    filter_no_select_video_paths = set()
    no_select_file_path = os.path.join(out_path, 'video.noselect.txt')
    if os.path.exists(no_select_file_path):
        with open(no_select_file_path, 'r') as file:
            for video_path in file.readlines():
                filter_no_select_video_paths.add(video_path.strip())

    file_counts = []
    with open(video_file_list, 'r') as file:
        for line in file.readlines():
            line_info = line.strip().split()
            file_path = line_info[0]
            face_count = int(line_info[2])
            file_counts.append((file_path, face_count))

    file_counts.sort(key=lambda x: x[1], reverse=True)

    video_paths = []
    face_counts = []
    for (file_path, face_count) in file_counts:
        if file_path in filter_select_video_paths:
            continue
        if file_path in filter_no_select_video_paths:
            continue
        video_paths.append(file_path)
        face_counts.append(face_count)

    out_file = open(select_file_path, 'a+')
    out_nosel_file = open(no_select_file_path, 'a+')
    try:
        total_video_count = len(video_paths)
        is_exit = False
        sleep_time = SLEEP_TIME
        for i, video_path in enumerate(video_paths):
            print('{}/{} {} {}'.format((i+1), total_video_count, video_path, face_counts[i]))

            full_video_path = os.path.join(base_path, video_path)
            cap = cv2.VideoCapture(full_video_path)
            total_frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            imWidth = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            imHeight = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            frameRate = int(cap.get(cv2.CAP_PROP_FPS))
            frame_count = 0
            force_pause = False
            sleep_time_str = ''
            while True:
                grabbed, img = cap.read()
                if not grabbed:
                    break

                frame_count += 1
                draw_frame(img, frame_count, total_frame_count)
                if (img.shape[0] != DISPLAY_HEIGHT) or (img.shape[1] != DISPLAY_WIDTH):
                    img = cv2.resize(img, (DISPLAY_WIDTH, DISPLAY_HEIGHT))
                cv2.imshow('image', img)

                pause = (total_frame_count-1) <= frame_count or force_pause

                if pause:
                    sleep_time_str = ''
                    while True:
                        image = img.copy()
                        if len(sleep_time_str) > 1:
                            cv2.putText(image, sleep_time_str, (64, 64), cv2.FONT_HERSHEY_PLAIN, 2.0, (0, 0, 255), thickness=2)
                        cv2.imshow('image', image)
                        key = cv2.waitKey()
                        if key == ord('n') or key == ord('N'):
                            if video_path not in filter_no_select_video_paths:
                                info_str = '{}\n'.format(video_path)
                                out_nosel_file.write(info_str)
                                print('No Select:', video_path)
                                filter_no_select_video_paths.add(video_path)
                            break
                        elif key == 27 or key == ord('q' or key == ord('Q')):
                            break
                        elif key == ord('s') or key == ord('S'):
                            if video_path not in filter_select_video_paths:
                                info_str = '{}\n'.format(video_path)
                                out_file.write(info_str)
                                print('Select:', video_path)
                                filter_select_video_paths.add(video_path)
                            break
                        elif key == ord('r') or key == ord('R'):  # replay
                            cap.release()
                            cap = cv2.VideoCapture(full_video_path)
                            frame_count = 0
                            force_pause = False
                            if sleep_time_str != '':
                                sleep_time = int(sleep_time_str)
                            break
                        elif ord('0') <= key <= ord('9'):  # sleep time
                            if sleep_time_str == '':
                                sleep_time_str = chr(key)
                            else:
                                sleep_time_str += chr(key)
                        elif key == ord('t') or key == ord('T'):  # set sleep time
                            if sleep_time_str != '':
                                sleep_time = int(sleep_time_str)
                                sleep_time_str = ''
                        elif key == ord('z') or key == ord('Z'):  # reset sleep time
                            sleep_time = SLEEP_TIME
                        # elif key == ord('p') or key == ord('P'):
                        #     force_pause = True
                        elif key == ord('c') or key == ord('C'):
                            force_pause = False
                            break
                else:
                    if 0 == sleep_time:
                        key = cv2.waitKey()
                    else:
                        key = cv2.waitKey(sleep_time)
                if key == ord('n') or key == ord('N'):
                    if video_path not in filter_no_select_video_paths:
                        info_str = '{}\n'.format(video_path)
                        out_nosel_file.write(info_str)
                        print('No Select:', video_path)
                        filter_no_select_video_paths.add(video_path)
                    break
                elif key == ord('s') or key == ord('S'):
                    if video_path not in filter_select_video_paths:
                        info_str = '{}\n'.format(video_path)
                        out_file.write(info_str)
                        print('Select:', video_path)
                        filter_select_video_paths.add(video_path)
                    break
                elif key == ord('r') or key == ord('R'):  # replay
                    cap.release()
                    cap = cv2.VideoCapture(full_video_path)
                    frame_count = 0
                    force_pause = False
                    if sleep_time_str != '':
                        sleep_time = int(sleep_time_str)
                        sleep_time_str = ''
                elif ord('0') <= key <= ord('9'):  # sleep time
                    if sleep_time_str == '':
                        sleep_time_str = chr(key)
                    else:
                        sleep_time_str += chr(key)
                elif key == ord('t') or key == ord('T'):  # set sleep time
                    if sleep_time_str != '':
                        sleep_time = int(sleep_time_str)
                        sleep_time_str = ''
                elif key == ord('z') or key == ord('Z'):  # reset sleep time
                    sleep_time = SLEEP_TIME
                elif key == ord('p') or key == ord('P'):
                    force_pause = True
                elif key == ord('c') or key == ord('C'):
                    force_pause = False
                elif key == ord(' '):
                    sleep_time = 0
                elif key == 27 or key == ord('q' or key == ord('Q')):
                    is_exit = True
                    break
            cap.release()
            if is_exit:
                break

        cv2.destroyAllWindows()
    except Exception as e:
        print(e)
    out_file.close()
    out_nosel_file.close()
    os.system('chmod a+wr {}'.format(select_file_path))
    os.system('chmod a+wr {}'.format(no_select_file_path))

    print('Finish!')


if __name__ == '__main__':
    main()
